This PR introduces the zen decorator, which changes the interface of an arbitrary (Hydra-agnostic) function so that it can accept a Hydra config as its input. The decorator inspects a inner-function's signature and extracts (+resolves and instantiates) the appropriate fields from an input config to call said function:

from hydra_zen import make_config, zen

def func(a, b, c):
    return a + b + c

>>> zen_wrapped_func = zen(func)
>>> Cfg = make_config(a=1, b=builds(int, 2), c="${a}", unused=100)
# extracts Cfg.a/b/c, resolves interpolated fields, and instantiates targeted configs, and calls `func(a, b, c)`
>>> zen_wrapped_func(Cfg)  #  a=1 + b=2 + c=1 -> 4
4

Using @zen to improve task functions

@zen is designed to help decouple one's task function from the Hydra framework. By doing so, it improves the task function's legibility, versatility, and testability.

Given:

from hydra_zen import make_config, builds, just, instantiate

class Foo:
    def __init__(self, x: str) -> None:
        pass

Cfg = make_config(seed=1, foo=builds(Foo, x="bar"), unused=[1, 2])

One typically writes task functions as so:

def old_task_fn(cfg):
    seed = cfg.seed  # manually accessing and instantiating attributes
    foo = instantiate(cfg.foo)
    print(seed, foo)

>>> old_task_fn(Cfg)
1 <__main__.Foo object at 0x00000262295F0D00>

There are several issues with this:

• The input – cfg – is relatively opaque to users and type-checkers alike.
• The process of manually accessing and instantiating attributes is tedious and produces boilerplate code (it isn't so bad for this toy example, but it gets bad quickly in practice)
• old_task_fn is tightly coupled to the Hydra framework – you must pass it a config to run it

@zen strives to rectify all of these shortcomings. Let's use it to refactor our task function:

from hydra_zen import zen

@zen
def new_task_fn(seed: int, foo: Foo):
    print(seed, foo)

>>> new_task_fn(Cfg)
1 <__main__.Foo object at 0x00000262295F0D00>

Here, @zen makes new_task_fn explicit, legible, and boilerplate free. It works by inspecting the signature of new_task_fn and extracting and instantiating the corresponding parameters from our config¹². Given the explicit signature (with optional annotations), users and IDEs can easily understand the context of the task function's body.

Furthermore, one can run the underlying task function, via .func, independently of a Hydra app:

>>> new_task_fn.func(seed=10, foo=Foo("hi"))
10 <__main__.Foo object at 0x00000262295F4790>

Given this accessibility, and because our task-function is now free of Hydra-specific boilerplate code, we can easily use/test our task function outside of the context of our Hydra app.

zen makes it trivial to take any 3rd party function and transform it into a Hydra-compatible task function:

Using @zen instead of @hydra.main

The object returned by zen provides a convenience method -- Zen.hydra_main -- so that users need not double-wrap with @hydra.main to create a CLI:

# example.py
from hydra.core.config_store import ConfigStore

from hydra_zen import builds, zen

def f(x: int, y: int):
    print(x, y)

cs = ConfigStore.instance()
cs.store(name="my_app", node=builds(f, populate_full_signature=True))


if __name__ == "__main__":
    zen(f).hydra_main(config_name="my_app", config_path=None)

$ python example.py x=1 y=2
1 2

Additional Bells & Whistles

Validation

A zen-wrapped function can validate configs without calling the function itself. This makes it easy to test compatibility between your task functions and configs (e.g., as part of your CI/CD process)

>>> def f(x: int): ...
>>> zen_f = zen(f)
>>> zen_f.validate({"x": 1})  # OK
>>> zen_f.validate({"y": 1})  # Missing x
---------------------------------------------------------------------------
HydraZenValidationError: `cfg` is missing the following fields: x

Customizing the Wrapper Behavior

One can subclass hydra_zen.wrapper.Zen and pass it to @zen to modify the wrapped behavior.
In the following example we add the ability to log the config (as a yaml) upon each call of a zen-wrapped function.

from hydra_zen.wrapper import Zen
from hydra_zen import to_yaml, zen

class MyZen(Zen):
    def __call__(self, __cfg):
        print(f"Logged:\n{to_yaml(__cfg)}")
        return super().__call__(__cfg)

>>> @zen(ZenWrapper=MyZen)
... def fn(x: int, y: str): return (x, y)

>>> fn(dict(x=1, y="hi"))
Logged:
x: 1
'y': hi

(1, 'hi')

Adding a Pre-Call Step

Recall that @zen will automatically instantiate a sub-config prior to passing it to the decorated function. If that instantiated object relies on random behavior, it can be useful to be able to set a seed prior to the instantiation process. We can do this via @zen(pre_call=...):

Without pre-call:

import random
from hydra_zen import builds, zen


@zen
def f(rand_val: int):
    return rand_val


Cfg = dict(rand_val=builds(random.randint, 0, 10))

>>> [f(Cfg) for _ in range(10)]
[3, 8, 2, 4, 2, 1, 9, 4, 8, 9]

With pre-call:

import random
from hydra_zen import builds, zen

# Note that we use zen to also wrap the pre-call function: to extract `seed` from the config
@zen(pre_call=zen(lambda seed: random.seed(seed)))  
def f(rand_val: int):
    return rand_val


Cfg = dict(
        rand_val=builds(random.randint, 0, 10),
        seed=0,
)

>>> [f(Cfg) for _ in range(10)]
[6, 6, 6, 6, 6, 6, 6, 6, 6, 6]

Validation propagates through zen-wrapped pre-call functions:

>>> def f(x: int): ...
>>> zen_f = zen(f, pre_call=zen(lambda seed: None))
>>> zen_f.validate({"x": 1, "seed": 10})  # OK
>>> zen_f.validate({"x": 1})  # Missing seed as required by pre-call 
---------------------------------------------------------------------------
HydraZenValidationError: `cfg` is missing the following fields: seed

Passing Through The Config

Some task functions require complete access to the full config to gain access to sub-configs. One can specify the field named zen_config³ in their task function's signature to signal zen that it should pass the full config to that parameter .

from hydra_zen import zen

def f(x: int, zen_cfg):
    return x, zen_cfg

>>> zen(f)(dict(x=1, y="${x}"))
(1, {'x': 1, 'y': 1})